"""
预测模型代理
负责机器学习模型的训练和预测
"""

from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error
import numpy as np
import joblib
import os

class PredictionModelAgent:
    """预测模型代理"""
    
    def __init__(self, model_dir='models'):
        """初始化预测模型"""
        self.models = {}
        self.model_dir = model_dir
        os.makedirs(self.model_dir, exist_ok=True)
    
    def train_predictive_model(self, X, y, model_type='random_forest', save_model=False):
        """
        训练预测模型
        参数:
            X: 特征数据
            y: 目标变量
            model_type: 模型类型(默认随机森林)
            save_model: 是否保存模型
        返回:
            训练好的模型和评估指标
        """
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
        
        if model_type == 'random_forest':
            model = RandomForestRegressor(n_estimators=100, random_state=42)
            model.fit(X_train, y_train)
            
            # 评估模型
            y_pred = model.predict(X_test)
            mse = mean_squared_error(y_test, y_pred)
            
            self.models[model_type] = model
            
            if save_model:
                self._save_model(model, model_type)
            
            return {
                'model': model,
                'metrics': {
                    'mse': mse,
                    'feature_importances': model.feature_importances_.tolist()
                }
            }
        
        raise ValueError(f"未知模型类型: {model_type}")

    def _save_model(self, model, model_name):
        """保存模型到文件"""
        path = os.path.join(self.model_dir, f"{model_name}.joblib")
        joblib.dump(model, path)
    
    def load_model(self, model_name):
        """从文件加载模型"""
        path = os.path.join(self.model_dir, f"{model_name}.joblib")
        if os.path.exists(path):
            model = joblib.load(path)
            self.models[model_name] = model
            return model
        return None